import os
import Image
import numpy as np


def generate_patches(patch_size=32, stride=14,
                     saved_path='/media/dell/cb552bf1-c649-4cca-8aca-3c24afca817b/dell/wxm/Data/SRCNNset/patches32_test'):
    train_pgms = []
    pgm_folder = '/media/dell/cb552bf1-c649-4cca-8aca-3c24afca817b/dell/wxm/Data/SRCNNset/Test'
    flag_pgm = 0
    x_one = 0
    y_one = 0
    x_two = patch_size
    y_two = patch_size
    # nb_patches = 1000
    for img in os.listdir(pgm_folder):
        patch_flag = 0
        flag_pgm += 1
        print 'nb.{0} images'.format(flag_pgm)
        train_pgms.append(img)
        img_path = pgm_folder + '/' + img
        img_obj = Image.open(img_path)
        for i in range(int((512 - patch_size) / stride + 1)):
            for j in range(int((512 - patch_size) / stride + 1)):
                patch_flag += 1
                box = (x_one + j * stride, y_one + i * stride, x_two + j * stride, y_two + i * stride)
                img_patch = img_obj.crop(box)
                img_patch.save(saved_path + '/' + str(patch_flag) + '_' + img)
                if patch_flag % 200 == 0:
                    print 'nb.{0} patches'.format(patch_flag)


def generate_patches_walk(patch_size=32, stride=14,
                          saved_path='/media/dell/cb552bf1-c649-4cca-8aca-3c24afca817b/dell/wxm/Data/SRCNNset/patches32_test'):
    train_pgms = []
    pgm_folder = '/media/dell/cb552bf1-c649-4cca-8aca-3c24afca817b/dell/wxm/Data/SRCNNset/Test'
    flag_pgm = 0
    x_one = 0
    y_one = 0
    x_two = patch_size
    y_two = patch_size
    # nb_patches = 1000
    for dirpath, dirname, filname in os.walk(pgm_folder):
        if 'Set5' in dirpath:
            saved_path = '/media/dell/cb552bf1-c649-4cca-8aca-3c24afca817b/dell/wxm/Data/SRCNNset/patches32_test/Set5'
            print 'Set5'
        elif 'Set14' in dirpath:
            print 'Set14'
            saved_path = '/media/dell/cb552bf1-c649-4cca-8aca-3c24afca817b/dell/wxm/Data/SRCNNset/patches32_test/Set14'
        else:
            continue
        for img in filname:
            patch_flag = 0
            flag_pgm += 1
            print 'nb.{0} images'.format(flag_pgm)
            train_pgms.append(img)
            img_path = dirpath + '/' + img
            img_obj = Image.open(img_path)
            for i in range(int((512 - patch_size) / stride + 1)):
                for j in range(int((512 - patch_size) / stride + 1)):
                    patch_flag += 1
                    box = (x_one + j * stride, y_one + i * stride, x_two + j * stride, y_two + i * stride)
                    img_patch = img_obj.crop(box)
                    img_patch.save(saved_path + '/' + str(patch_flag) + '_' + img)
                    if patch_flag % 200 == 0:
                        print 'nb.{0} patches'.format(patch_flag)

if __name__ == '__main__':
    generate_patches_walk()
